import argparse
from functools import cached_property
import logging
import os
from typing import Any, Generator
import gzip
from pathlib import Path
from pydantic import BaseModel, model_validator
import polars as pl
import pytaxonkit


logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("GCGenome")


class PathogenGenomics(BaseModel):
    gcpathogen_dir: Path
    taxdump_dir: Path
    outdir: Path

    @model_validator(mode="after")
    def validate_outdir(self):
        self.outdir.mkdir(parents=True, exist_ok=True)
        return self

    @cached_property
    def reference_gca_txt(self) -> Path:
        return self.gcpathogen_dir / "pathogen_referenceGCA.txt"

    @cached_property
    def genome_fasta(self) -> Path:
        return self.outdir / "pathogen.genome.fasta"

    @cached_property
    def genome_faidx(self) -> Path:
        return self.genome_fasta.with_suffix(".fasta.fai")

    @cached_property
    def summary_tsv(self) -> Path:
        return self.outdir / "pathogen.summary.tsv"

    @cached_property
    def species_tsv(self) -> Path:
        return self.outdir / "pathogen.species.tsv"

    @cached_property
    def accession_tsv(self) -> Path:
        return self.outdir / "pathogen.accession.tsv"

    def write_species(self) -> pl.DataFrame:
        logger.info("read %s", self.reference_gca_txt)
        records = pl.read_csv(self.reference_gca_txt, separator="\t").cast(pl.String)
        records = records.rename({"taxonid": "taxid"})
        logger.info("run TaxonKit")
        genus_records = (
            pl.DataFrame(pl.from_pandas(pytaxonkit.lineage(records["taxid"].unique().to_list(), data_dir=self.taxdump_dir)))
            .cast(pl.String)
            .rename({"TaxID": "taxid", "Name": "name", "Rank": "rank"})
            .with_columns(pl.col("LineageTaxIDs").str.split(";").alias("lineage_taxids"))
            .with_columns(pl.col("Lineage").str.split(";").alias("lineages"))
            .with_columns(pl.col("lineage_taxids").list.get(-1).replace("", None).fill_null(pl.col("taxid")).alias("species_taxid"))
            .with_columns(pl.col("lineage_taxids").list.get(-2).replace("", "-").alias("genus_taxid"))
            .with_columns(pl.col("lineages").list.get(-1).replace("", None).fill_null(pl.col("name")).alias("species_name"))
            .with_columns(pl.col("lineages").list.get(-2).replace("", "-").alias("genus_name"))
            .select(["taxid", "name", "species_taxid", "species_name", "genus_taxid", "genus_name", "rank"])
        )
        records = (
            records.join(genus_records, on="taxid", how="left")
            .fill_null("-")
            .select(["taxid", "name", "species_taxid", "species_name", "genus_taxid", "genus_name", "organism_type", "ref", "rank"])
        )
        logger.info("write %s", self.species_tsv)
        records.write_csv(self.species_tsv, separator="\t")

    def iter_fna(self) -> Generator[Path, Any, None]:
        organism_types = pl.read_csv(self.species_tsv, separator="\t", columns=["organism_type"])["organism_type"].unique()
        for fna_dir in map(lambda x: self.gcpathogen_dir / x, organism_types):
            yield from fna_dir.glob("*.fna.gz")

    def write_genome_fasta(self):
        rows = []
        with open(self.genome_fasta, "w") as writer:
            for fna_file in self.iter_fna():
                ref = fna_file.name.replace(".fna.gz", "")
                with gzip.open(fna_file, "rt") as reader:
                    for line in reader:
                        if line[0] == ">":
                            accession = line[1:].split(" ", 1)[0]
                            rows.append([ref, accession])
                        writer.write(line)
                logger.info("written %s to %s", fna_file, self.genome_fasta)
        pl.DataFrame(rows, schema=["ref", "accession"]).write_csv(self.accession_tsv, separator="\t")

    def write_summary(self):
        logger.info("write %s", self.accession_tsv)
        species_records = pl.read_csv(self.species_tsv, separator="\t", schema_overrides={"genus_taxid": str, "species": str})
        accession_records = pl.read_csv(self.accession_tsv, separator="\t").cast(pl.String)
        records = accession_records.join(species_records, on="ref", how="left").select(
            ["accession", "taxid", "name", "species_taxid", "species_name", "genus_taxid", "genus_name", "organism_type", "rank", "ref"]
        )
        records.write_csv(self.summary_tsv, separator="\t")

    def build_faidx(self):
        command = f"samtools faidx {self.genome_fasta}"
        logger.info(command)
        if os.system(command) > 0:
            raise RuntimeError(f"failed to run {command}")

    def run(self):
        if not self.species_tsv.exists():
            self.write_species()
        # if not (self.species_tsv.exists() and self.genome_fasta.exists()):
        #     self.write_genome_fasta()
        if not self.summary_tsv.exists():
            self.write_summary()
        # if not self.genome_faidx.exists():
        #     self.build_faidx()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="gcGenome")
    parser.add_argument("-g", "--gcpathogen_dir", type=Path, default=Path("/mnt/nas/database/pathogen_2/gcpathogen/"), help="data directory")
    parser.add_argument("-t", "--taxdump_dir", type=Path, default=Path("/home/zhuy/Database/ncbi/taxdump/"), help="taxdump directory")
    parser.add_argument("-o", "--outdir", type=Path, default=Path("/mnt/nas/database/pathogen_2/pathogen_genomics/"), help="output directory")
    parser.set_defaults(func=lambda args: PathogenGenomics.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
